import copy
import numpy as np
from OCC.Core.BRep import BRep_Builder
from OCC.Core.BRepBuilderAPI import BRepBuilderAPI_Transform
from OCC.Core.TopoDS import TopoDS_Compound
from OCC.Core.TopoDS import TopoDS_Shape
from OCC.Core.gp import gp_Trsf
from trimesh.transformations import identity_matrix, concatenate_matrices

try:
    from Symppart import SymPart
except:
    from .Symppart import SymPart

try:
    from group import Group,Body, Part, PART_DIR, GROUP_DIR
except:
    from .group import Group,Body, Part, PART_DIR, GROUP_DIR

import sympy as sym

from dopy.transform import vector_rotation_matrix
from trimesh.transformations import translation_matrix, unit_vector

__version__ = "1.7.0"
print('AIAD 3D Parser', __version__)


def get_extra_args(args, k):
    arr = []
    for name in args:
        ss = name.split('__', 1)
        if ss[0] == k:
            dic = {}
            dic.update(args[name])
            dic['name'] = ss[1]

            arr.append(dic)

    return arr


class SymGroup(SymPart):
    def __init__(self, group_type,
                 group_dir=GROUP_DIR,
                 part_dir=PART_DIR, extra_args=None, global_cache=None
                 ):
        # if global_cache is not None:
        #     Global_Cache.Cache=copy.deepcopy(global_cache)
        self.parse_element_to_sym(group_type, group_dir, prefix='.')

        self.parse_extra_args_sym(extra_args)

        for k in self.elements:
            element = self.elements[k]
            category = element['category']

            sub_extra_args = get_extra_args(self.args, k)

            if category.startswith('.'):
                element['instance'] = SymGroup(category[1:], group_dir, part_dir, extra_args=sub_extra_args)
            else:
                element['instance'] = SymPart(category, part_dir, extra_args=sub_extra_args)

        self.get_transform_sym()

    def get_child(self, name):
        ss = name.split('/')
        if len(ss) == 1:
            return self.elements[name]['instance'], self.elements[name]['transform_sym']
        else:
            raise NotImplementedError()

    def get_transform_sym(self):
        for ele in self.trans_sym.values():
            if 'align' in ele:
                a0 = [
                    self.get_point_sym(ele['align'][0] + '.+x'),
                    self.get_point_sym(ele['align'][0] + '.+y'),
                    self.get_point_sym(ele['align'][0] + '.+z')
                ]

                a1 = [
                    self.get_point_sym(ele['align'][1] + '.+x'),
                    self.get_point_sym(ele['align'][1] + '.+y'),
                    self.get_point_sym(ele['align'][1] + '.+z')
                ]

                a0 = a0 - self.get_point_sym(ele['align'][0] + '.o')
                a1 = a1 - self.get_point_sym(ele['align'][1] + '.o')

                m = vector_rotation_matrix(a0, a1)
                self.add_transform(ele['group'], m)

            if 'rotate' in ele:
                axis_map = {
                    "+x": Part.AXIS_X,
                    "+y": Part.AXIS_Y,
                    "+z": Part.AXIS_Z,
                }
                axis = axis_map[ele['rotate'][1].lower()]
                angle = sym.pi / 180 * ele['rotate'][0]
                sina = sym.sin(angle)
                cosa = sym.cos(angle)
                direction = axis[1]
                point = axis[0]
                direction = unit_vector(direction[:3])
                # rotation matrix around unit vector
                M = np.diag([cosa, cosa, cosa, 1.0])
                M[:3, :3] += np.outer(direction, direction) * (1.0 - cosa)

                direction = direction * sina
                M[:3, :3] += np.array([[0.0, -direction[2], direction[1]],
                                       [direction[2], 0.0, -direction[0]],
                                       [-direction[1], direction[0], 0.0]])

                # if point is specified, rotation is not around origin
                if point is not None:
                    point = np.array(point[:3], dtype=np.float64, copy=False)
                    M[:3, 3] = point - np.dot(M[:3, :3], point)

                m = sym.Matrix(M)

                self.add_transform_sym(ele['group'], m)

            if 'point2point' in ele:
                d1 = self.get_point_sym(ele['point2point'][0])
                d2 = self.get_point_sym(ele['point2point'][1])

                v = np.array(d2) - np.array(d1)
                # print(v)

                if len(ele['point2point']) >= 3:
                    keep = ele['point2point']
                    if 'x' not in keep:
                        v[0] = 0
                    if 'y' not in keep:
                        v[1] = 0
                    if 'z' not in keep:
                        v[2] = 0

                m = sym.Matrix([[1, 0, 0, v[0]], [0, 1, 0, v[1]], [0, 0, 1, v[2]], [0, 0, 0, 1]])

                self.add_transform_sym(ele['group'], m)

            if 'move' in ele:
                v = ele['move']
                m = sym.Matrix([[1, 0, 0, v[0]], [0, 1, 0, v[1]], [0, 0, 1, v[2]], [0, 0, 0, 1]])

                self.add_transform_sym(ele['group'], m)

    def get_point_sym(self, name):
        if isinstance(name, list):
            assert len(name) == 3
            return np.array(name, dtype=float)

        if '/' not in name:
            assert name in self.points

            refer = self.points[name].get('refer')
            offset = self.points[name].get('offset')

            if refer is None:
                coor = sym.Matrix([[0, 0, 0]])
            else:
                assert '/' in refer

                coor = SymGroup.get_point_sym(self, refer)

            if offset is not None:
                coor += np.array([offset])
        else:
            ss = name.split('/')
            assert len(ss) > 1
            # element = self
            m = sym.Matrix([[1, 0, 0, 0], [0, 1, 0, 0], [0, 0, 1, 0], [0, 0, 0, 1]])
            for e in ss[:-1]:
                element, tran = self.get_child(e)
                m = m * tran

            if isinstance(element, SymPart):
                coor = element.get_point_sym(ss[-1])
                for args in element.syms:
                    outargs = ss[0] + '__' + args
                    if outargs not in self.syms:
                        self.syms[outargs] = sym.symbols(outargs)
                    coor = [e.subs({sym.symbols(args): self.syms[outargs]}) for e in coor]

            else:
                coor = element.get_point_sym(ss[-1])
            coor = sym.Matrix([[coor[0], coor[1], coor[2], 1]])

            coor = m * coor.T

        return coor[0:3]

    def get_point_finaly_sym(self, name):
        coor = self.get_point_sym(name)
        for arg in self.syms:
            if arg not in self.args:
                first_index = arg.find('__')
                element_name = arg[:first_index]
                element_arg = arg[first_index + 2:]
                element = self.elements[element_name]['instance']
                coor = [e.subs({self.syms[arg]: element.args[element_arg]['v']}) for e in coor]
        return coor

    def get_point(self, name):
        coor = self.get_point_sym(name)
        for arg in self.syms:
            if arg not in self.args:
                first_index = arg.find('__')
                element_name = arg[:first_index]
                element_arg = arg[first_index + 2:]
                element = self.elements[element_name]['instance']
                coor = [e.subs({self.syms[arg]: element.args[element_arg]['v']}) for e in coor]
            else:
                coor = [e.subs({self.syms[arg]: self.args[arg]['v']}) for e in coor]
        return coor

    def shape(self, debug=False):
        mesh_all = None
        name_all =[]
        for name in self.elements:
            name_all.append(name)
            element = self.elements[name]
            m = element['instance'].compound()
            a = [[element['transform'][i][j] for j in range(4)] for i in range(3)]

            T = gp_Trsf()
            myshape = TopoDS_Shape(m)
            T.SetValues(a[0][0], a[0][1], a[0][2], a[0][3],
                        a[1][0], a[1][1], a[1][2], a[1][3],
                        a[2][0], a[2][1], a[2][2], a[2][3])
            trans = BRepBuilderAPI_Transform(T)
            trans.Perform(myshape)
            new_shape = trans.Shape()
            if mesh_all is None:
                mesh_all = [new_shape]
            else:
                mesh_all.append(new_shape)
        return mesh_all,name_all

    def compound(self):
        myNewShape = self.shape()
        if isinstance(myNewShape,list):
            compound = TopoDS_Compound()
            B = BRep_Builder()
            B.MakeCompound(compound)
            for shape in myNewShape:
                B.Add(compound, shape)

            return compound

        else:
            compound = TopoDS_Compound()
            B = BRep_Builder()
            B.MakeCompound(compound)
            B.Add(compound,myNewShape)
            return compound


class SymBody(SymGroup):
    def __init__(self, group_type,
                 group_dir=GROUP_DIR,
                 part_dir=PART_DIR, extra_args=None,
                 transform=None, global_cache=None
                 ):
        self.transform = transform
        if self.transform is None:
            self.transform = sym.Matrix(identity_matrix())

        super().__init__(group_type, group_dir=group_dir, part_dir=part_dir, extra_args=extra_args,
                         global_cache=global_cache)

    def apply_transform(self, m):
        self.transform = sym.Matrix(concatenate_matrices(m, self.transform))

    def apply_translation(self, v):
        self.transform = sym.Matrix(concatenate_matrices(translation_matrix(v), self.transform))

    def get_point(self, name):
        coor = super().get_point_finaly_sym(name)
        coor = sym.Matrix([[coor[0], coor[1], coor[2], 1]])
        coor = self.transform * coor.T
        return coor[0:3]

    def shape(self, debug=False):
        mesh_all = super().shape(debug)[0]
        myNewShape = mesh_all
        compound = TopoDS_Compound()
        B = BRep_Builder()
        B.MakeCompound(compound)
        for shape in myNewShape:
            B.Add(compound, shape)
        a = [[self.transform[i][j] for j in range(4)] for i in range(3)]

        T = gp_Trsf()
        myshape = TopoDS_Shape(compound)
        T.SetValues(a[0][0], a[0][1], a[0][2], a[0][3],
                    a[1][0], a[1][1], a[1][2], a[1][3],
                    a[2][0], a[2][1], a[2][2], a[2][3])
        trans = BRepBuilderAPI_Transform(T)
        trans.Perform(myshape)
        new_shape = trans.Shape()

        return new_shape


if __name__ == '__main__':

    x = SymGroup('clamp_4a')
    print(x.get_point_finaly_sym('Topleft/A.-z'))
    possible_para_dict = copy.deepcopy(x.args)
    possible_para_sym = [x.syms[e] for e in possible_para_dict]
    # possible_para_sym_dict=[x.syms[] for e in possible_para_dict]
    # f=sym.lambdify((L1,L2,SIZE,H,ARC), test, 'numpy')

    x1 = Group('clamp_4a')
    # args=[
    #     {
    #         'name': "上压紧块__H",
    #         'v': 30
    #     },
    #     {
    #         'name': "上压紧块__L1",
    #         'v': 20,
    #         'range': [20,200]
    #     },
    #             {
    #         'name': "上压紧块__L2",
    #         'v': 50
    #     },
    #     {
    #         'name': "上压紧块__ARC",
    #         'v': 180
    #     },
    # ]
    # test=x.get_point_sym('上压紧块/A.-z')

    coor = x.get_point_finaly_sym('Topleft/A.-z')
    for arg in x.args:
        coor = [e.subs({x.syms[arg]: x.args[arg]['v']}) for e in coor]
    print(coor)
    print(x1.get_point('Topleft/A.-z'))

    # test=[sym.N(e.subs({SIZE:20,L1:80,L2:50,H:25,ARC:105})) for e in test]
    # print(test)
    # for arg in args:
    #     test=[e.subs({x.syms[arg['name']]:arg['v']}) for e in test]
    # test=[sym.N(e) for e in test]

    # x = Body('clamp_1b')
    # import detail_design.jxq.model_parameter_update as model_parameter_update
    # import detail_design.jxq.model_helper as model_helper

    # print(model_parameter_update.get_clamp_point(x))
    # # model_helper.get_clamp_point_normal(x)
    # angle=80/180*math.pi
    # # rotate_matrix = trimesh.transformations.rotation_matrix(angle, [0,0,1], [0,0,0])
    #
    # # x.apply_transform(rotate_matrix)
    # # x.apply_translation([1000, 0, 0])
    # # x.show()
    # x.show([x.get_point('setting_pt')])
    # # x.show()
    # #
    # sys.exit()

    # x = Group('c1', extra_args=[{'v': 500, 'name': 'L1__width'}])
    # x.show()

    # x = Body("clamp_1f")
    # x.show()
    # # x.apply_translation([100, 0, 0])
    # # x.show(x.get_point("c1/L1/A.+x+y+z"), radius=5)
    # sys.exit()

    # x = Group("clamp_2c_ex")
    # # start=time.time()
    # # for i in range(1000):
    # #     x = Group("clamp_2c_ex")
    # #     # with open("clamp.json5", "r", encoding="utf-8") as fp:
    # #     #     obj = json5.load(fp)
    # #     # pprint(x.args)
    # # time_save = time.time()-start
    # #
    # # pprint(time_save)
    #
    # # print(time.time() - time_save)
    #
    # # print(x.get_point('上压紧块/A.-z+x'))
    # print(x.get_point('RPS-1-AB'))
    # # x.get_mesh().show()
    # x.show([x.get_point('RPS-1-AB'), x.get_point('下压紧块/A.-z-x+y')])
    # pprint(x.obj)
    # x.get_mesh().show()
    # x.show()
